# -*- coding: utf-8 -*-
import os.path
import sys
import time

import numpy as np
from PyQt5.QtGui import (
    QBrush,
    QPainter,
    QPen,
    QPixmap,
    QKeySequence,
    QPen,
    QBrush,
    QColor,
    QImage,
)
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import (
    QFileDialog,
    QApplication,
    QGraphicsEllipseItem,
    QGraphicsItem,
    QGraphicsRectItem,
    QGraphicsScene,
    QGraphicsView,
    QGraphicsPixmapItem,
    QHBoxLayout,
    QPushButton,
    QSlider,
    QVBoxLayout,
    QWidget,
    QShortcut,
    QCheckBox,
    QLabel,
    QSpinBox,
    QSplitter
)

from numpy import random, zeros, array, uint8, repeat, sum, clip
from skimage import transform, io
from torch import manual_seed, cuda, device, no_grad, as_tensor, tensor, float, sigmoid
from torch.nn import functional as F
from PIL import Image, ImageDraw
from segment_anything import sam_model_registry

# freeze seeds
manual_seed(2023)
cuda.empty_cache()
cuda.manual_seed(2023)
random.seed(2023)

SAM_MODEL_TYPE = "vit_b"
MedSAM_CKPT_PATH = "work_dir/MedSAM/medsam_vit_b.pth"
MEDSAM_IMG_INPUT_SIZE = 1024
device = device("cuda")


@no_grad()
def medsam_inference(medsam_model, img_embed, box_1024, height, width):
    box_torch = as_tensor(box_1024, dtype=float, device=img_embed.device)
    if len(box_torch.shape) == 2:
        box_torch = box_torch[:, None, :]  # (B, 1, 4)

    sparse_embeddings, dense_embeddings = medsam_model.prompt_encoder(
        points=None,
        boxes=box_torch,
        masks=None,
    )
    low_res_logits, _ = medsam_model.mask_decoder(
        image_embeddings=img_embed,  # (B, 256, 64, 64)
        image_pe=medsam_model.prompt_encoder.get_dense_pe(),  # (1, 256, 64, 64)
        sparse_prompt_embeddings=sparse_embeddings,  # (B, 2, 256)
        dense_prompt_embeddings=dense_embeddings,  # (B, 256, 64, 64)
        multimask_output=False,
    )

    low_res_pred = sigmoid(low_res_logits)  # (1, 1, 256, 256)

    low_res_pred = F.interpolate(
        low_res_pred,
        size=(height, width),
        mode="bilinear",
        align_corners=False,
    )  # (1, 1, gt.shape)
    low_res_pred = low_res_pred.squeeze().cpu().numpy()  # (256, 256)
    medsam_seg = (low_res_pred > 0.5).astype(uint8)
    return medsam_seg


print("Loading MedSAM model, a sec.")
tic = time.perf_counter()

# set up model
medsam_model = sam_model_registry["vit_b"](checkpoint=MedSAM_CKPT_PATH).to(device)
medsam_model.eval()

print(f"Done, took {time.perf_counter() - tic}")


def np2pixmap(np_img):
    height, width, channel = np_img.shape
    bytesPerLine = 3 * width
    qImg = QImage(np_img.data, width, height, bytesPerLine, QImage.Format_RGB888)
    return QPixmap.fromImage(qImg)


def pixmap2np(pixmap):
    return array(pixmap.toImage())


colors = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
    (0, 255, 255),
    (128, 0, 0),
    (0, 128, 0),
    (0, 0, 128),
    (128, 128, 0),
    (128, 0, 128),
    (0, 128, 128),
    (255, 255, 255),
    (192, 192, 192),
    (64, 64, 64),
    (255, 0, 255),
    (0, 255, 255),
    (255, 255, 0),
    (0, 0, 127),
    (192, 0, 192),
]


class Window(QWidget):
    def __init__(self):
        super().__init__()

        # configs
        self.half_point_size = 5  # radius of bbox starting and ending points

        # app stats
        self.image_path = None
        self.color_idx = 0
        self.bg_img = None
        self.is_mouse_down = False
        self.rect = None
        self.point_size = self.half_point_size * 2
        self.start_point = None
        self.end_point = None
        self.start_pos = (None, None)
        self.embedding = None
        self.prev_mask = None
        self.image_index = 0
        self.image_path_list = []

        self.EraserMode = False
        self.PainterMode = False

        self.view = QGraphicsView()
        self.view.setRenderHint(QPainter.Antialiasing)

        self.load_image_dir()

        # pixmap = self.load_image()

        # load_button = QPushButton("选取文件夹")
        # save_button = QPushButton("保存图片")
        #
        # hbox = QHBoxLayout(self)
        #
        # hbox.addWidget(load_button)
        # hbox.addWidget(save_button)

        # load_button.clicked.connect(self.load_image_dir)
        # save_button.clicked.connect(self.save_mask)

        # keyboard shortcuts
        self.quit_shortcut = QShortcut(QKeySequence("Ctrl+Q"), self)
        self.quit_shortcut.activated.connect(lambda: quit())

        self.undo_shortcut = QShortcut(QKeySequence("Ctrl+Z"), self)
        self.undo_shortcut.activated.connect(self.undo)

        vbox = QVBoxLayout(self)
        vbox.addWidget(self.view)

        # 新建垂直子布局用于放置按键
        sub_layout = QVBoxLayout()

        # 设置主布局内边距以及控件间距为10px
        sub_layout.setSpacing(10)
        # 设置此子布局和内部控件的间距为10px
        sub_layout.setContentsMargins(10, 10, 10, 10)

        self.__cbtn_Eraser = QCheckBox(" 使用橡皮擦 ")
        self.__cbtn_Eraser.setParent(self)
        self.__cbtn_Eraser.clicked.connect(self.on_cbtn_Eraser_clicked)
        sub_layout.addWidget(self.__cbtn_Eraser)

        self.__cbtn_Painter = QCheckBox(" 使用画笔 ")
        self.__cbtn_Painter.setParent(self)
        self.__cbtn_Painter.clicked.connect(self.on_cbtn_Painter_clicked)
        sub_layout.addWidget(self.__cbtn_Painter)

        splitter = QSplitter(self)  # 占位符
        sub_layout.addWidget(splitter)

        self.__label_penThickness = QLabel(self)
        self.__label_penThickness.setText("画笔粗细")
        self.__label_penThickness.setFixedHeight(20)
        sub_layout.addWidget(self.__label_penThickness)

        self.__spinBox_penThickness = QSpinBox(self)
        self.__spinBox_penThickness.setMaximum(20)
        self.__spinBox_penThickness.setMinimum(1)
        self.__thickness = 5
        self.__spinBox_penThickness.setValue(5)  # 默认粗细为10
        self.__spinBox_penThickness.setSingleStep(1)  # 最小变化值为2
        self.__spinBox_penThickness.valueChanged.connect(
            self.on_PenThicknessChange)  # 关联spinBox值变化信号和函数on_PenThicknessChange

        sub_layout.addWidget(self.__spinBox_penThickness)

        # vbox.addLayout(hbox)
        vbox.addLayout(sub_layout)
        self.setLayout(vbox)

    def undo(self):
        if self.prev_mask is None:
            self.setWindowTitle('回退失败，已无先前标注')
            print("No previous mask record")
            return

        self.color_idx -= 1

        bg = Image.fromarray(self.img_3c.astype("uint8"), "RGB")
        mask = Image.fromarray(self.prev_mask.astype("uint8"), "RGB")
        img = Image.blend(bg, mask, 0.4)

        self.scene.removeItem(self.bg_img)
        self.bg_img = self.scene.addPixmap(np2pixmap(array(img)))

        self.mask_c = self.prev_mask
        if sum(self.mask_c) ==0:
            self.prev_mask = zeros((*self.img_3c.shape[:2], 3), dtype="uint8")
        else:
            self.prev_mask = None

    # override the key press event
    def keyPressEvent(self, event):
        # if left arrow key is pressed
        if event.key() == Qt.Key_Q:
            if sum(self.mask_c) > 0:
                self.save_mask()
            self.image_index -= 1
            self.load_image()
            print("上一张图片切换完成")
        # if down arrow key is pressed
        elif event.key() == Qt.Key_E:
            if sum(self.mask_c) > 0:
                self.save_mask()
            self.image_index += 1
            if self.image_index >= len(self.image_path_list):
                print("已经是最后一张图片，关闭")
            self.load_image()
            print("下一张图片切换完成")

        elif event.key() == Qt.Key_Space:
            self.save_mask()

    def load_image_dir(self):
        directory = QFileDialog.getExistingDirectory(self, "选取文件夹", "./")  # 起始路径

        self.image_path_list = [os.path.join(directory, file) for file in sorted(os.listdir(directory)) if
                                (file[-3:] == 'png' or file[-3:] == 'jpg' or file[-3:] == 'bmp') and file[
                                                                                                     -8:-4] != 'mask']

        mask_list = [os.path.join(directory, file) for file in sorted(os.listdir(directory)) if file[-8:-4] == 'mask']
        if len(mask_list) > 0:
            last_mask_name = mask_list[-1].replace('_mask', "")
            self.image_index = self.image_path_list.index(last_mask_name)
        print(self.image_path_list)

        self.load_image()

    def load_image(self):

        file_path = self.image_path_list[self.image_index]
        img_np = io.imread(file_path)
        if len(img_np.shape) == 2:
            img_3c = repeat(img_np[:, :, None], 3, axis=-1)
        else:
            img_3c = img_np

        self.img_3c = img_3c
        self.image_path = file_path
        self.get_embeddings()
        pixmap = np2pixmap(self.img_3c)

        H, W, _ = self.img_3c.shape

        self.scene = QGraphicsScene(0, 0, W, H)
        self.end_point = None
        self.rect = None
        self.bg_img = self.scene.addPixmap(pixmap)
        self.bg_img.setPos(0, 0)
        self.mask_c = zeros((*self.img_3c.shape[:2], 3), dtype="uint8")
        self.prev_mask = zeros((*self.img_3c.shape[:2], 3), dtype="uint8")
        self.view.setScene(self.scene)

        # events
        self.scene.mousePressEvent = self.mouse_press
        self.scene.mouseMoveEvent = self.mouse_move
        self.scene.mouseReleaseEvent = self.mouse_release

    def mouse_press(self, ev):
        if self.PainterMode or self.EraserMode:
            # 鼠标按下时，获取鼠标的当前位置保存为上一次位置
            self.__currentPos = (int(ev.scenePos().x()), int(ev.scenePos().y()))
            self.__lastPos = self.__currentPos
        else:
            x, y = ev.scenePos().x(), ev.scenePos().y()
            self.is_mouse_down = True
            self.start_pos = ev.scenePos().x(), ev.scenePos().y()
            self.start_point = self.scene.addEllipse(
                x - self.half_point_size,
                y - self.half_point_size,
                self.point_size,
                self.point_size,
                pen=QPen(QColor("red")),
                brush=QBrush(QColor("red")),
            )

    def mouse_move(self, ev):
        if self.PainterMode or self.EraserMode:
            # 鼠标移动时，更新当前位置，并在上一个位置和当前位置间画线
            self.__currentPos = (int(ev.scenePos().x()), int(ev.scenePos().y()))

            # self.prev_mask = self.mask_c.copy()
            img_mask = Image.fromarray(self.mask_c.astype("uint8"), "RGB")
            img_mask_draw = ImageDraw.Draw(img_mask)
            if self.PainterMode:
                # 绘制
                img_mask_draw.line([self.__lastPos,self.__currentPos], fill=colors[self.color_idx % len(colors)],width=self.__thickness)
            else:
                # 橡皮擦模式下画笔为纯黑色，粗细为10
                img_mask_draw.line([self.__lastPos,self.__currentPos], fill='black', width=self.__thickness)
            self.__lastPos = self.__currentPos
            self.mask_c = np.array(img_mask).astype("uint8")
            bg = Image.fromarray(self.img_3c, "RGB")
            mask = Image.fromarray(self.mask_c.astype("uint8"), "RGB")
            img = Image.blend(bg, mask, 0.4)
            self.scene.removeItem(self.bg_img)
            self.bg_img = self.scene.addPixmap(np2pixmap(array(img)))

        else:
            if not self.is_mouse_down:
                return

            x, y = ev.scenePos().x(), ev.scenePos().y()

            if self.end_point is not None:
                self.scene.removeItem(self.end_point)
            self.end_point = self.scene.addEllipse(
                x - self.half_point_size,
                y - self.half_point_size,
                self.point_size,
                self.point_size,
                pen=QPen(QColor("red")),
                brush=QBrush(QColor("red")),
            )

            if self.rect is not None:
                self.scene.removeItem(self.rect)
            sx, sy = self.start_pos
            xmin = min(x, sx)
            xmax = max(x, sx)
            ymin = min(y, sy)
            ymax = max(y, sy)
            self.rect = self.scene.addRect(
                xmin, ymin, xmax - xmin, ymax - ymin, pen=QPen(QColor("red"))
            )

    def mouse_release(self, ev):
        if not self.EraserMode and not self.PainterMode:
            x, y = ev.scenePos().x(), ev.scenePos().y()
            sx, sy = self.start_pos
            xmin = min(x, sx)
            xmax = max(x, sx)
            ymin = min(y, sy)
            ymax = max(y, sy)

            self.is_mouse_down = False

            H, W, _ = self.img_3c.shape
            box_np = array([[xmin, ymin, xmax, ymax]])
            # print("bounding box:", box_np)
            box_1024 = box_np / array([W, H, W, H]) * 1024

            sam_mask = medsam_inference(medsam_model, self.embedding, box_1024, H, W)

            self.prev_mask = self.mask_c.copy()
            self.color_idx += 1
            self.mask_c[sam_mask != 0] = colors[self.color_idx % len(colors)]


            bg = Image.fromarray(self.img_3c.astype("uint8"), "RGB")
            mask = Image.fromarray(self.mask_c.astype("uint8"), "RGB")
            img = Image.blend(bg, mask, 0.4)

            self.scene.removeItem(self.bg_img)
            self.bg_img = self.scene.addPixmap(np2pixmap(array(img)))

    def save_mask(self):
        out_path = f"{self.image_path.split('.')[0]}_mask.png"
        io.imsave(out_path, self.mask_c)
        self.setWindowTitle('图片保存成功')

    @no_grad()
    def get_embeddings(self):
        print("加载中")
        self.setWindowTitle('加载中')
        img_1024 = transform.resize(
            self.img_3c, (1024, 1024), order=3, preserve_range=True, anti_aliasing=True
        ).astype(uint8)
        img_1024 = (img_1024 - img_1024.min()) / clip(
            img_1024.max() - img_1024.min(), a_min=1e-8, a_max=None
        )  # normalize to [0, 1], (H, W, 3)
        # convert the shape to (3, H, W)
        img_1024_tensor = (
            tensor(img_1024).float().permute(2, 0, 1).unsqueeze(0).to(device)
        )

        # if self.embedding is None:
        with no_grad():
            self.embedding = medsam_model.image_encoder(
                img_1024_tensor
            )  # (1, 256, 64, 64)
        self.setWindowTitle('加载成功')

    def on_cbtn_Eraser_clicked(self):
        if self.__cbtn_Eraser.isChecked():
            self.EraserMode = True  # 进入橡皮擦模式
            self.PainterMode = False
            self.__cbtn_Painter.setChecked(False)
        else:
            self.EraserMode = False  # 退出橡皮擦模式

    def on_cbtn_Painter_clicked(self):
        if self.__cbtn_Painter.isChecked():
            self.PainterMode = True  # 进入橡皮擦模式
            self.EraserMode = False
            self.__cbtn_Eraser.setChecked(False)
        else:
            self.PainterMode = False  # 退出橡皮擦模式

    def on_PenThicknessChange(self):
        penThickness = self.__spinBox_penThickness.value()
        self.ChangePenThickness(penThickness)

    def ChangePenThickness(self, thickness=10):
        # 改变画笔粗细
        self.__thickness = thickness


app = QApplication(sys.argv)

w = Window()
w.show()

app.exec()
